import akshare as ak
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
from sklearn.preprocessing import MinMaxScaler
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
import os
from datetime import datetime, timedelta

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

print("=== 茅台股票20天价格预测系统（简化版）===")

def calculate_rsi(prices, window=14):
    """计算RSI相对强弱指数"""
    delta = prices.diff()
    gain = (delta.where(delta > 0, 0)).rolling(window=window).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window=window).mean()
    rs = gain / loss
    rsi = 100 - (100 / (1 + rs))
    return rsi

def get_maotai_data():
    """获取茅台股票数据"""
    try:
        print("正在获取茅台股票数据...")
        stock_code = "600519"
        
        # 获取过去2年的数据
        end_date = datetime.now().strftime('%Y%m%d')
        start_date = (datetime.now() - timedelta(days=730)).strftime('%Y%m%d')
        
        df = ak.stock_zh_a_hist(symbol=stock_code, period="daily", start_date=start_date, end_date=end_date, adjust="qfq")
        
        if len(df) == 0:
            raise Exception("获取到空数据")
            
        # 重命名列
        df = df.rename(columns={
            '日期': 'date',
            '开盘': 'open',
            '收盘': 'close', 
            '最高': 'high',
            '最低': 'low',
            '成交量': 'volume'
        })
        
        df['date'] = pd.to_datetime(df['date'])
        df = df.set_index('date')
        
        print(f"成功获取 {len(df)} 条真实数据")
        return df, "真实"
        
    except Exception as e:
        print(f"获取真实数据失败: {e}")
        print("使用模拟数据...")
        return create_mock_data(), "模拟"

def create_mock_data():
    """创建模拟数据"""
    np.random.seed(42)
    
    end_date = datetime.now()
    start_date = end_date - timedelta(days=730)
    dates = pd.date_range(start=start_date, end=end_date, freq='D')
    dates = [d for d in dates if d.weekday() < 5]  # 只保留工作日
    
    base_price = 1650
    n_days = len(dates)
    
    # 生成价格数据
    price_changes = np.random.normal(0, 20, n_days)
    prices = [base_price]
    
    for change in price_changes[1:]:
        new_price = prices[-1] + change
        new_price = max(1200, min(2200, new_price))
        prices.append(new_price)
    
    # 生成其他数据
    opens, highs, lows, closes, volumes = [], [], [], prices, []
    
    for i, close in enumerate(closes):
        if i == 0:
            open_price = close + np.random.normal(0, 5)
        else:
            open_price = closes[i-1] + np.random.normal(0, 10)
        
        daily_range = abs(np.random.normal(0, 30))
        high = max(open_price, close) + daily_range * 0.7
        low = min(open_price, close) - daily_range * 0.3
        volume = max(100000, int(np.random.normal(500000, 200000)))
        
        opens.append(open_price)
        highs.append(high)
        lows.append(low)
        volumes.append(volume)
    
    df = pd.DataFrame({
        'open': opens,
        'high': highs,
        'low': lows,
        'close': closes,
        'volume': volumes
    }, index=dates)
    
    return df

# 获取数据
df, data_type = get_maotai_data()

# 创建数据目录
if not os.path.exists('data'):
    os.makedirs('data')

# 保存数据
if data_type == "真实":
    df.to_csv('data/maotai_real_data.csv')
    print("真实数据已保存")

print(f"使用{data_type}数据，范围: {df.index[0].strftime('%Y-%m-%d')} 到 {df.index[-1].strftime('%Y-%m-%d')}")

# 数据预处理
print("开始数据预处理...")

# 创建技术指标
df['ma5'] = df['close'].rolling(window=5).mean()
df['ma20'] = df['close'].rolling(window=20).mean()
df['rsi'] = calculate_rsi(df['close'])
df['previous_close'] = df['close'].shift(1)
df['price_change'] = df['close'] - df['previous_close']

# 删除NaN
df = df.dropna()

print(f"预处理后数据量: {len(df)}")

# 显示基本信息
current_price = df['close'].iloc[-1]
print(f"\\n当前股价: ¥{current_price:.2f}")
print(f"5日均线: ¥{df['ma5'].iloc[-1]:.2f}")
print(f"20日均线: ¥{df['ma20'].iloc[-1]:.2f}")
print(f"RSI: {df['rsi'].iloc[-1]:.1f}")

# 绘制历史数据
plt.figure(figsize=(15, 8))

plt.subplot(2, 1, 1)
plt.plot(df.index, df['close'], linewidth=2, color='red', label='收盘价')
plt.plot(df.index, df['ma5'], linewidth=1, color='blue', alpha=0.7, label='5日均线')
plt.plot(df.index, df['ma20'], linewidth=1, color='green', alpha=0.7, label='20日均线')
plt.title(f'茅台股票价格走势（{data_type}数据）')
plt.ylabel('价格（元）')
plt.legend()
plt.grid(True, alpha=0.3)

plt.subplot(2, 1, 2)
plt.plot(df.index, df['rsi'], color='purple')
plt.axhline(y=70, color='r', linestyle='--', alpha=0.7, label='超买线')
plt.axhline(y=30, color='g', linestyle='--', alpha=0.7, label='超卖线')
plt.title('RSI指标')
plt.ylabel('RSI')
plt.legend()
plt.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

# 准备机器学习数据
print("准备机器学习数据...")

# 选择特征
features = ['open', 'high', 'low', 'volume', 'ma5', 'ma20', 'rsi', 'previous_close']
X = df[features].values
y = df['close'].values

# 数据归一化
scaler_X = MinMaxScaler()
scaler_y = MinMaxScaler()

X_scaled = scaler_X.fit_transform(X)
y_scaled = scaler_y.fit_transform(y.reshape(-1, 1)).flatten()

# 创建时序数据（使用过去10天预测下一天）
def create_sequences(X, y, seq_len=10):
    X_seq, y_seq = [], []
    for i in range(seq_len, len(X)):
        X_seq.append(X[i-seq_len:i])
        y_seq.append(y[i])
    return np.array(X_seq), np.array(y_seq)

seq_length = 10
X_seq, y_seq = create_sequences(X_scaled, y_scaled, seq_length)

print(f"序列数据形状: X={X_seq.shape}, y={y_seq.shape}")

# 分割训练测试数据
train_size = int(len(X_seq) * 0.8)
X_train = X_seq[:train_size]
X_test = X_seq[train_size:]
y_train = y_seq[:train_size]
y_test = y_seq[train_size:]

# 为线性回归重塑数据
X_train_flat = X_train.reshape(X_train.shape[0], -1)
X_test_flat = X_test.reshape(X_test.shape[0], -1)

# 训练线性回归模型
print("训练线性回归模型...")
model = LinearRegression()
model.fit(X_train_flat, y_train)

# 预测
y_pred_test = model.predict(X_test_flat)

# 反归一化
y_test_real = scaler_y.inverse_transform(y_test.reshape(-1, 1)).flatten()
y_pred_test_real = scaler_y.inverse_transform(y_pred_test.reshape(-1, 1)).flatten()

# 评估模型
mse = mean_squared_error(y_test_real, y_pred_test_real)
mae = mean_absolute_error(y_test_real, y_pred_test_real)
r2 = r2_score(y_test_real, y_pred_test_real)

print(f"\\n模型评估:")
print(f"均方误差: {mse:.2f}")
print(f"平均绝对误差: {mae:.2f}")
print(f"R²得分: {r2:.4f}")

# 预测未来20天
print("\\n预测未来20天股价...")

# 获取最后seq_length天的数据
last_sequence = X_scaled[-seq_length:]
future_predictions = []

for i in range(20):
    # 预测下一天
    next_pred = model.predict(last_sequence.reshape(1, -1))[0]
    future_predictions.append(next_pred)
    
    # 更新序列
    # 简化处理：使用预测值更新部分特征
    new_row = last_sequence[-1].copy()
    new_row[0] = next_pred  # 开盘价 ≈ 前收盘价
    new_row[1] = next_pred * 1.02  # 最高价
    new_row[2] = next_pred * 0.98  # 最低价
    new_row[-1] = next_pred  # previous_close
    
    last_sequence = np.vstack([last_sequence[1:], new_row])

# 反归一化预测结果
future_predictions_real = scaler_y.inverse_transform(np.array(future_predictions).reshape(-1, 1)).flatten()

# 生成未来日期
last_date = df.index[-1]
future_dates = []
for i in range(1, 21):
    next_date = last_date + timedelta(days=i)
    # 跳过周末
    while next_date.weekday() >= 5:
        next_date += timedelta(days=1)
    future_dates.append(next_date)

# 可视化预测结果
plt.figure(figsize=(15, 8))

# 历史数据（最近60天）+ 未来预测
recent_days = 60
recent_dates = df.index[-recent_days:]
recent_prices = df['close'][-recent_days:].values

plt.plot(recent_dates, recent_prices, label='历史价格', color='blue', linewidth=2)
plt.plot(future_dates, future_predictions_real, label='未来20天预测', color='red', linewidth=2, marker='o')

# 连接线
plt.plot([recent_dates[-1], future_dates[0]], [recent_prices[-1], future_predictions_real[0]], 
         color='orange', linewidth=2, linestyle='--', alpha=0.7)

plt.title('茅台股票未来20天价格预测')
plt.xlabel('日期')
plt.ylabel('股价（元）')
plt.legend()
plt.grid(True, alpha=0.3)
plt.xticks(rotation=45)
plt.tight_layout()
plt.show()

# 输出预测结果
print(f"\\n=== 预测结果摘要 ===")
print(f"当前价格: ¥{current_price:.2f}")
print(f"预测第1天: ¥{future_predictions_real[0]:.2f}")
print(f"预测第20天: ¥{future_predictions_real[-1]:.2f}")

change_1day = future_predictions_real[0] - current_price
change_20day = future_predictions_real[-1] - current_price
change_1day_pct = (change_1day / current_price) * 100
change_20day_pct = (change_20day / current_price) * 100

print(f"1天预期变化: {change_1day:+.2f} ({change_1day_pct:+.2f}%)")
print(f"20天预期变化: {change_20day:+.2f} ({change_20day_pct:+.2f}%)")

print(f"\\n=== 详细预测（前10天）===")
for i in range(10):
    change = future_predictions_real[i] - current_price
    change_pct = (change / current_price) * 100
    print(f"第{i+1:2d}天 ({future_dates[i].strftime('%m-%d')}): ¥{future_predictions_real[i]:7.2f} "
          f"({change:+6.2f}, {change_pct:+5.2f}%)")

# 保存预测结果
prediction_df = pd.DataFrame({
    'date': future_dates,
    'predicted_price': future_predictions_real,
    'change_amount': future_predictions_real - current_price,
    'change_percent': ((future_predictions_real - current_price) / current_price) * 100
})

prediction_df.to_csv('data/maotai_20days_prediction.csv', index=False)
print(f"\\n预测结果已保存到 data/maotai_20days_prediction.csv")

# 技术分析建议
avg_prediction = np.mean(future_predictions_real)
if avg_prediction > current_price * 1.05:
    suggestion = "看涨 📈"
elif avg_prediction < current_price * 0.95:
    suggestion = "看跌 📉"
else:
    suggestion = "震荡 ➡️"

print(f"\\n=== 投资建议 ===")
print(f"基于模型预测: {suggestion}")
print(f"预测平均价格: ¥{avg_prediction:.2f}")
print(f"\\n⚠️  风险提示: 股票预测仅供参考，投资需谨慎！")

print(f"\\n✅ 茅台股票20天预测完成！")